import sqlite3


class DBSet:
    def __init__(self, db_path: str):
        self.db_path = db_path
        self.conn = sqlite3.connect(self.db_path)
        self.cursor = self.conn.cursor()
        # Update table creation to include local_url

    def set(self, md5: str, media_id: str, media_url: str, local_url: str = "") -> None:
        """Set cache entry for file"""
        # Update REPLACE INTO statement to include local_url
        self.cursor.execute(
            "REPLACE INTO cache (md5, media_id, media_url, local_url) VALUES (?, ?, ?, ?)",
            (md5, media_id, media_url, local_url),
        )
        self.conn.commit()

    def get(self, md5: str) -> list:
        """Get cache entry for file"""
        # Update SELECT statement to include local_url
        self.cursor.execute(
            "SELECT md5, media_id, media_url, local_url FROM cache WHERE md5 = ?",
            (md5,),
        )
        result = self.cursor.fetchone()
        if result is None:
            return []
        md5 = result[0]
        media_id = result[1]
        media_url = result[2]
        local_url = result[3]
        return [md5, media_id, media_url, local_url]

    def all_md5s_in_cache(self) -> list:
        """Get all md5s in cache"""
        self.cursor.execute("SELECT md5 FROM cache")
        results = self.cursor.fetchall()
        return [result[0] for result in results]

    def close(self):
        """Close the database connection"""
        self.conn.close()
